package cn.bytepulse.pojo;


import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;

import java.io.*;

/**
 * @author jiejiebiezheyang
 * @since 2024-02-02 19:00
 */
public class CachedResponseWrapper extends HttpServletResponseWrapper {

    private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    private PrintWriter writer;

    public CachedResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return new ServletOutputStream() {
            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener listener) {
                throw new UnsupportedOperationException();
            }

            @Override
            public void write(int b) throws IOException {
                outputStream.write(b);
            }
        };
    }

    @Override
    public PrintWriter getWriter() throws UnsupportedEncodingException {
        if (writer == null) {
            writer = new PrintWriter(new OutputStreamWriter(outputStream, getCharacterEncoding()));
        }
        return writer;
    }

    /**
     * 获取缓存的响应数据（字节数组）
     */
    public byte[] getCachedData() {
        if (writer != null) {
            writer.flush();
        }
        return outputStream.toByteArray();
    }

    /**
     * 获取缓存的响应数据（字符串）
     */
    public String getCachedBody() throws UnsupportedEncodingException {
        return new String(getCachedData(), getCharacterEncoding());
    }
}